import json
import logging
import pandas as pd
import re
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import JsonOutputParser

import os

from catalog.dip_catalog import DipCatalog
from series.agent.bed_supplementary_catalog import BedSupplementaryCatalog
from series.rules.dip_match_tcm_superior_disease_rules import MatchTcmSuperiorDiseaseGroupingBase
from series.rules.dip_synthetic_rules import DIPSyntheticRules
from util.read_txt_file import ReadTxt
from series.agent.age_supplementary_catalog import AgeSupplementaryCatalog
from series.dao.prompt_knowledge_base_dao import PromptKnowledgeBaseDao

logger = logging.getLogger(__name__)


class DIPGroupRules:
    def __init__(self):
        logger.info("开始初始化大模型")
        self.db = PromptKnowledgeBaseDao()
        # 从环境变量读取配置
        load_dotenv()
        api_key = os.getenv("OPENAI_API_KEY")
        base_url = os.getenv("OPENAI_BASE_URL")
        model_name = os.getenv("MODEL_NAME")
        logger.info("初始化DIP分组规则系统")
        self.bed_agent = BedSupplementaryCatalog()

        # 初始化年龄智能体
        self.age_agent = AgeSupplementaryCatalog()

        # 初始化大模型
        self.llm = ChatOpenAI(
            temperature=0.1,
            max_tokens=300,
            base_url=base_url,
            model=model_name,
            api_key=api_key
        )
        logger.info("大模型初始化完成")

        # # 获取数据库中的提示词知识库
        # sql_request = self.db.query_base_usr()
        # logger.info(f"获取到 {len(sql_request)} 条提示词知识库数据")
        #
        # # 处理提示词数据
        # prompt_knowledge = ""
        # if sql_request:
        #     prompt_knowledge = "\n".join([
        #         item.get('cue_word_content', '')
        #         for item in sql_request
        #     ])
        #     logger.info("提示词知识库加载完成")
        # else:
        #     logger.warning("未获取到提示词知识库数据")
        #
        # # 存储提示词知识库以供后续使用
        # self.prompt_knowledge = prompt_knowledge
        #
        # self.query_prompt = ChatPromptTemplate.from_template(
        #     prompt_knowledge
        # )

        self.query_prompt = ChatPromptTemplate.from_template(
            """你是DIP病种分组专家，请根据DIP病种目录库数据进行精确分组。

患者信息：
- 诊断编码：{diagnosis_code}
- 诊断名称：{diagnosis_name}
- 患者年龄：{patient_age}岁
- 住院天数：{hospitalization_days}天
- 总费用: {total_cost}
- 中医费用 :{tcm_cost}
- 主手术名称：{main_surgery_code}
- 次要手术名称：{other_surgery_code}
- 治疗操作：{treatment_code}
- 诊断操作：{diagnostic_operation}
提示必须诊断名称和手术加在一起进行匹配
核心逻辑：一定要匹配最多记录的记录 比如主手术+次手术 主手术+次手术+治疗操作 都匹配上 一定要取主手术+次手术+治疗操作
1. 判断患者是否有手术（主手术/次手术/治疗操作/诊断操作是否全部为空）
   - 无手术 → 核心保守治疗
   - 有手术 → 按优先级匹配复合编码：
     主手术 > 次手术 > 治疗性操作 > 诊断性操作
2. 复合编码解析：
   - "+" 表示 AND（必须同时具备）
   - "/" 表示 OR（任意一个即可）
3. 匹配失败执行综合目录调整
4. 床日调整：住院天数>21时执行
5. 年龄调整：床日已执行则跳过
6. 中医优势病种调整：床日已执行则跳过，中医费用占比>=60%时执行

DIP目录库匹配数据：
{catalog_data}

统一入组规则：
{unified_rules}

请根据目录库数据，返回JSON格式的结果：

{{
  "分组名称": "匹配成功无手术情况显示[核心保守],有手术显示核心手术",
  "诊断编码": "只返回编码数值，如Z51.8，不包含名称和括号",
  "基础分组说明"："分组详细说明"
  "分值": 数值类型,
  "手术及操作编码": "匹配成功的手术及操作编码",
  "需要综合目录分析": "核心目录匹配失败需要调整综合目录true/false",
  "需要床日调整": true/false,
  "床日调整原因": "调整原因或无",
  "需要年龄调整": true/false :年龄调整范围 0-28天,29天-1岁,1-6岁,7-17岁,66-120岁"
  "年龄调整原因": "年龄调整范围或无",
  "需要中医优势病种调整": true/false,
  "中医优势病种调整原因": "中医优势病种调整原因或无",
  "是否匹配成功": "核心目录匹配上true匹配不上false",
  "是否有手术": "手术及操作编码不等于空的时候true等于空false"
}}

**重要要求**：
1. 必须返回有效的JSON格式

返回JSON："""
        )

        # 创建处理链
        self.json_parser = JsonOutputParser()
        self.chain = self.query_prompt | self.llm | self.json_parser

    def merge_all_results(self, base_result, synthetic_result, bed_result, age_result):
        """简单合并所有结果"""
        final_result = base_result.copy()

        # 合并综合目录结果
        if synthetic_result and isinstance(synthetic_result, dict):
            if synthetic_result.get('综合调整完成'):
                final_result['分值'] = synthetic_result.get('分值', final_result.get('分值'))
                final_result['综合调整完成'] = True
                final_result['综合调整说明'] = synthetic_result.get('调整说明', '')

        # 合并床日结果
        if bed_result and isinstance(bed_result, dict):
            if bed_result.get('床日调整完成'):
                final_result['分值'] = bed_result.get('分值', final_result.get('分值'))
                final_result['床日调整完成'] = True
                final_result['床日调整说明'] = bed_result.get('调整说明', '')

        # 合并年龄结果
        if age_result and isinstance(age_result, dict):
            if age_result.get('年龄调整完成'):
                final_result['分值'] = age_result.get('分值', final_result.get('分值'))
                final_result['年龄调整完成'] = True
                final_result['年龄调整说明'] = age_result.get('调整说明', '')

        return final_result

    def parse_ai_result(self, ai_result):
        """处理AI返回的JSON结果"""
        # 如果AI返回的已经是字典格式，直接返回
        if isinstance(ai_result, dict):
            return ai_result

        # 如果是字符串，尝试解析JSON
        if isinstance(ai_result, str):
            try:
                return json.loads(ai_result)
            except json.JSONDecodeError:
                logger.error(f"AI返回的不是有效JSON格式: {ai_result}")
                return {
                    "分组名称": "",
                    "诊断编码": "",
                    "分值": None,
                    "手术及操作编码": "",
                    "需要床日调整": False,
                    "床日调整原因": "",
                    "需要年龄调整": False,
                    "年龄调整原因": "",
                    "需要中医优势病种调整": False,
                    "中医优势病种调整原因": "",
                    "是否有手术": "",
                    "是否匹配成功": "",
                    "错误": ""
                }

        # 其他情况返回错误结果
        return {
            "分组名称": "未知格式",
            "分值": None,
            "错误": f"未知的返回格式: {type(ai_result)}"
        }

    # 执行ai
    def execute_ai(self, patient_row, catalog_data):
        unified_rules = self.read_knowledge_base_txt()
        try:
            print(catalog_data)
            res = json.loads(catalog_data)
            first_record = res["data"][0]

            # 执行基本的DIP分组
            result = self.chain.invoke({
                "patient_info": patient_row,
                "prompt_knowledge": self.query_prompt,
                "diagnosis_code": patient_row['诊断编码'],
                "diagnosis_name": first_record['诊断名称'],
                'hospitalization_days': patient_row['住院天数'],
                "patient_age": patient_row.get('年龄', 0),
                "total_cost": patient_row['总费用'],
                "tcm_cost": patient_row['中医费用'],
                "main_surgery_code": patient_row.get('主手术', '') if patient_row.get('主手术') not in [None, '','nan'] else "",
                "other_surgery_code": patient_row.get('其它手术', '') if patient_row.get('其它手术') not in [None, '','nan'] else "",
                "treatment_code": patient_row.get('治疗操作', '') if patient_row.get('治疗操作') not in [None, '','nan'] else "",
                "diagnostic_operation": patient_row.get('诊断操作', '') if patient_row.get('诊断操作') not in [None, '','nan'] else "",
                "catalog_data": catalog_data,
                "unified_rules": unified_rules[:2200]
            })
            logger.info("基础大模型分析完成 {}".format(json.dumps(result, ensure_ascii=False)))
            # 解析AI结果(1)
            parsed_result = self.parse_ai_result(result)

            # 添加基础DIP分组分析原因
            synthetic_result = None
            synthetic_rules = DIPSyntheticRules()
            if synthetic_rules.ai_check_synthetic_agent(parsed_result, patient_row['诊断编码']):
                synthetic_result = synthetic_rules.call_synthetic_agent(patient_row, parsed_result)
                print("综合目录调整结果:", synthetic_result)

                # 立即合并综合目录结果
                if synthetic_result and isinstance(synthetic_result, dict) and synthetic_result.get('综合目录调整完成'):
                    parsed_result['分组名称'] = '综合手术'
                    parsed_result['分值'] = synthetic_result.get('分值', parsed_result.get('分值'))
                    parsed_result['综合调整完成'] = True
                    parsed_result['综合调整说明'] = synthetic_result.get('调整说明', '')
                    parsed_result['是否匹配成功'] = True  # 综合目录匹配成功
                    logger.info(f"综合目录调整完成: {parsed_result.get('分组名称')}, 分值: {parsed_result.get('分值')}")

            print("基础DIP分组结果:", parsed_result)
            bed_result = None
            if self.ai_check_bed_agent(parsed_result, patient_row['诊断编码']):
                bed_result = self.call_bed_agent(patient_row, parsed_result)
                # 立即合并床日结果
                if bed_result and isinstance(bed_result, dict) and bed_result.get('床日调整完成'):
                    parsed_result['分值'] = bed_result.get('分值', parsed_result.get('分值'))
                    parsed_result['床日调整完成'] = True
                    parsed_result['床日调整说明'] = bed_result.get('调整说明', '')
                    logger.info(f"床日调整后分值更新为: {parsed_result.get('分值')}")
            tcm_result = None
            tcm =MatchTcmSuperiorDiseaseGroupingBase()
            if tcm.ai_check_tcm_agent(parsed_result, patient_row['诊断编码']):
                tcm_result = tcm.call_tcm_agent(patient_row, result)
                print("中医优势病种调整后的结果:", tcm_result)
                # 立即合并床日结果
                if tcm_result and isinstance(tcm_result, dict) and tcm_result.get('中医优势病种调整完成'):
                    parsed_result['分值'] = tcm_result.get('分值', parsed_result.get('分值'))
                    parsed_result['中医优势病种调整完成'] = True
                    parsed_result['中医优势病种调整说明'] = tcm_result.get('调整说明', '')
                    logger.info(f"中医优势病种调整后分值更新为: {parsed_result.get('分值')}")
            # 让AI判断是否需要调用年龄智能体
            age_result = None
            if self.ai_check_age_agent(parsed_result, patient_row['诊断编码']):
                age_result = self.call_age_agent(patient_row, result)
                print("年龄调整后的结果:", age_result)

                # 立即合并年龄结果
                if isinstance(age_result, dict) and age_result.get('年龄调整完成'):
                    adjusted_score = age_result.get('分值')
                    if adjusted_score is not None:
                        parsed_result['分值'] = adjusted_score
                        parsed_result['年龄调整完成'] = True
                        parsed_result['需要年龄调整'] ='是'
                        parsed_result['年龄调整说明'] = age_result.get('调整说明', '')
                        parsed_result['分组名称'] = age_result.get('分组名称', '')
                        logger.info(f"年龄调整后分值更新为: {adjusted_score}")
                        # 更新result字符串以包含年龄调整信息
                        if isinstance(result, str):
                            result = result + f"\n\n【年龄智能体调整】\n{age_result.get('调整说明', '')}"
                        else:
                            result_str = json.dumps(result, ensure_ascii=False, indent=2) if isinstance(result,dict) else str(result)
                            result = result_str + f"\n\n【年龄智能体调整】\n{age_result.get('调整说明', '')}"

            # 返回原始文本和解析结果（已经在过程中合并了）
            return {
                "原始文本": result.strip() if isinstance(result, str) else str(result),
                "解析结果": parsed_result
            }

        except Exception as e:
            logger.error(f"执行AI时发生错误: {str(e)}")
            return {"错误": str(e)}

    # 检查是否执行年龄辅助目录
    def ai_check_age_agent(self, ai_result, diagnosis_code):
        dipCatalog = DipCatalog()
        catalog_age = dipCatalog.load_age_supplementary_catalog("诊断编码", diagnosis_code)
        try:
            age_res = json.loads(catalog_age)
        except Exception as e:
            logger.error(f"年龄辅助目录 JSON 解析失败: {e}, 原始数据: {catalog_age}")
            return False
        age_record = age_res.get("data", None)

        try:
            if not age_record:
                return False
            # 检查解析后的结果中是否需要年龄调整
            if ai_result.get('需要年龄调整', False):
                return True

            return True
        except Exception as e:
            logger.error(f"AI判断年龄智能体需求失败: {str(e)}")
            return False

    # 判断是有床日
    def ai_check_bed_agent(self, ai_result, diagnosis_code):
        dipCatalog = DipCatalog()
        catalog_bed = dipCatalog.load_bed_supplementary_catalog("诊断编码", diagnosis_code)
        # 解析 JSON
        try:
            bed_res = json.loads(catalog_bed)
        except Exception as e:
            logger.error(f"床日补助目录 JSON 解析失败: {e}, 原始数据: {catalog_bed}")
            return False

        # 安全获取 data
        bed_record = bed_res.get("data", None)

        try:
            if bed_record is None:
                return False
            # 检查解析后的结果中是否需要床日调整
            if ai_result.get('需要床日调整', False):
                return True
            return False

        except Exception as e:
            logger.error(f"AI判断床日智能体需求失败: {str(e)}")
            return False

    # 执行床日智能体
    def call_bed_agent(self, patient_row, base_result):
        """调用床日智能体"""
        try:
            logger.info(f"调用床日智能体，住院天数: {patient_row['住院天数']}天")
            dipCatalog = DipCatalog()
            diagnosis_code = base_result.get('诊断编码', '')
            catalog_bed = dipCatalog.load_bed_supplementary_catalog("诊断编码", diagnosis_code)

            bed_res = json.loads(catalog_bed)

            bed_catalog_data = bed_res["data"][0]

            # 调用床日智能体，现在返回JSON格式
            bed_result = self.bed_agent.analyze_bed_adjustment(patient_row, base_result, bed_catalog_data)

            # 直接返回JSON结果
            return bed_result

        except Exception as e:
            logger.error(f"床日智能体调用失败: {str(e)}")
            return {"错误": str(e)}

    # 执行年龄智能体
    def call_age_agent(self, patient_row, base_result):
        """调用年龄智能体"""
        try:
            logger.info(f"调用年龄智能体，患者年龄: {patient_row['年龄']}岁")
            dipCatalog = DipCatalog()
            catalog_age = dipCatalog.load_age_supplementary_catalog("诊断编码", patient_row['诊断编码'])
            # 调用年龄智能体，现在返回JSON格式
            age_result = self.age_agent.analyze_age_adjustment(patient_row, base_result, catalog_age)

            # 直接返回JSON结果
            return age_result

        except Exception as e:
            logger.error(f"年龄智能体调用失败: {str(e)}")
            return {"错误": str(e)}

    # 读取知识库(txt)
    def read_knowledge_base_txt(self):
        reader_txt = ReadTxt('规则文件.txt')
        return reader_txt.load_txt_lines()